#define BUFFERSIZE 1024
#define RECVSIZE 1040
#include <fstream>
#include <iostream>
#include <string.h>
#include <stdio.h>
#include <getopt.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdlib.h>
#include <netinet/tcp.h>
#include <fstream>
#include <iostream>
#include <botan/botan.h>
#include <malloc.h>
#include "RSA.h"
using namespace Botan;
using namespace std;

byte * encryption(RSA *rsa, PK_Encrypting_Key* enkey, byte *text, int text_len, u32bit *len)
{
  SecureVector<byte> ciphertext;
  rsa->encrypt(text,text_len,enkey,&ciphertext);

  Pipe pipe;
  pipe.process_msg(ciphertext);

  *len = pipe.remaining(0);

  byte *hand_shake_buf;
  hand_shake_buf = (byte *)malloc(*len);

  pipe.read(hand_shake_buf, *len);

  return(hand_shake_buf);
}

byte * decryption(RSA *rsa, PK_Decrypting_Key* dekey, byte *hand_shake_buf, int recv_size, size_t size)
{
  X509_PublicKey* ppkey=X509::load_key("rsapub1.pem");
  if(!ppkey)
       cout<<"no ppkey"<<endl;
  PK_Decrypting_Key* key = dynamic_cast<PK_Decrypting_Key*>(ppkey);
  if(!key)
       cout<<"no key"<<endl;
  SecureVector<byte> plaintext;
  PK_Decryptor* dec = get_pk_decryptor(*key, "EME1(SHA-1)");
                 try {
                 // will also accept a SecureVector<byte> as input
                         plaintext = dec->decrypt(hand_shake_buf,recv_size);
                 }
                 catch(Decoding_Error)
                 {
                 /* the ciphertext was invalid */
                 }

//cout<<"buffer: "<<(byte *)hand_shake_buf<<endl;
//cout<<"buffer size: "<<recv_size<<endl;
//  SecureVector<byte> ciphertext(hand_shake_buf,recv_size);
//cout<<"read into vector"<<endl;

// rsa->decrypt(ciphertext,*dekey,&plaintext);
cout<<"decrypt over"<<endl;

  Pipe pipe;
  pipe.process_msg(plaintext);

  recv_size = pipe.remaining(0);

  free(hand_shake_buf);
  hand_shake_buf = (byte *)malloc(recv_size);
  memset(hand_shake_buf,0,recv_size);
  pipe.read(hand_shake_buf,size);

  return(hand_shake_buf);
}

int main(int argc, char **argv)
{
  char *get_file = (char*) "sendfile.txt", *save_file = (char*) "receivefile.txt";
  char ch, *server_ip = (char *) "127.0.0.1";
  int port = 3344;

  while ((ch = getopt(argc, argv, "a:g:s:")) != -1) {
    switch (ch) {
      case 'g':
        get_file = optarg;
        break;
      case 's':
        save_file = optarg;
        break;
      case 'a':
        server_ip = optarg;
        break;
      case '?':
        if (optopt == 's' || optopt == 'g' || optopt == 'a') {
          fprintf(stderr, "Error: option -%c requires an argument\n", optopt);
        } else {
          fprintf(stderr, "Error: unknown option\n");
        }
        exit(0);
    }
  }

  struct sockaddr_in server_addr;
  int sock;

  if ((sock = socket(AF_INET,SOCK_STREAM, 0)) == -1) {
    fprintf(stderr, "Error: cannot open socket\n");
    exit(EXIT_FAILURE);
  }

  memset((char *) &server_addr, 0, sizeof(server_addr));
  server_addr.sin_family = AF_INET;
  server_addr.sin_port = htons(port);
  if (inet_aton(server_ip, &server_addr.sin_addr) == 0) {
    fprintf(stderr, "Error: invalid server address\n");
    exit(EXIT_FAILURE);
  }

  if(connect(sock,(struct sockaddr *)&server_addr,sizeof(server_addr)) < 0) {
    fprintf(stderr, "Error: cannot connect to server\n");
    exit(EXIT_FAILURE);
  }

  char recv_buf[RECVSIZE];
  fstream filestr;
  unsigned file_size;

  filestr.open (save_file, ios::out|ios::binary);
  LibraryInitializer init;
  string algo="AES";
  string passphrase="you never know";

  // text to encrypt
  byte symkey[]="I am symkey";
  byte autkey[]="I am authkey";
  u32bit len;
  RSA* rsa = new RSA();
  rsa->loadPublicKey("rsapub.pem");

  PK_Encrypting_Key* enkey = dynamic_cast<PK_Encrypting_Key*>(rsa->publicKey);
  if(!enkey) {
    cout << "no key" << endl;
  }

  //Handshake With Server
  int text_len = sizeof(symkey);
  byte *hand_shake_buf = encryption(rsa, enkey, symkey, text_len, &len);
  send(sock,hand_shake_buf,len,0);
  free(hand_shake_buf);
  text_len = sizeof(autkey);
  hand_shake_buf = encryption(rsa, enkey, autkey, text_len, &len);
  send(sock,hand_shake_buf,len,0);
  free(hand_shake_buf);
  text_len = strlen(get_file)+1;
  hand_shake_buf = encryption(rsa, enkey, (byte *)get_file, text_len, &len);
  send(sock,hand_shake_buf,len,0);
  free(hand_shake_buf);

  //AES SETUP --
  if(!have_block_cipher(algo)) {
    std::cout << "Don't know about the block cipher \"" << algo << "\"\n";
    return 1;
  }

  const u32bit key_len = max_keylength_of(algo);
  const u32bit iv_len = block_size_of(algo);

  std::auto_ptr<S2K> s2k(get_s2k("PBKDF2(SHA-1)"));
  s2k->set_iterations(8192);
  cout<<"Symmetric key: "<<symkey<<endl<<"Authentication key: " <<autkey<<endl;
  SymmetricKey bc_key = s2k->derive_key(key_len, (char*)symkey);//"BLK" + passphrase);
  InitializationVector iv = s2k->derive_key(iv_len,  (char *)symkey);//"IVL" + passphrase);
  SymmetricKey mac_key = s2k->derive_key(16, (char *)autkey);//"MAC" + passphrase);

  int r=0;
  r=recv(sock,recv_buf,16,0);
  //cout<<"Frist receive size"<<r<<endl;
  Pipe pipe(get_cipher(algo + "/CBC/PKCS7", bc_key, iv, DECRYPTION));
  pipe.start_msg();
  pipe.write((byte*)recv_buf,16);
  pipe.end_msg();
  pipe.read((byte*)recv_buf,4,0);
  file_size=*(int*)&recv_buf;
  printf("\nReady to receive file: %s\nFile size: %d bytes\n\n", get_file,file_size);

  unsigned recv_size = 0, write_size = 0, cum_send_size = 0;
  while  (cum_send_size != file_size) {
    char recv_mac[20]={0,};
    char our_mac[20]={0,};
    recv_size=recv(sock, recv_buf, RECVSIZE, 0);
    Pipe pipe(get_cipher(algo + "/CBC/PKCS7", bc_key, iv, DECRYPTION),
              new Fork(
                       0,
                       new MAC_Filter("HMAC(SHA-1)", mac_key)
                       )
             );
    write_size = (file_size-cum_send_size <= recv_size)?(file_size-cum_send_size):BUFFERSIZE;
    printf("Saved in all: %6d bytes,  Saved this time: %4d bytes\n",cum_send_size,write_size);
    pipe.start_msg();
    pipe.write((byte*)recv_buf,recv_size);
    pipe.end_msg();
    pipe.read((byte *)recv_buf,BUFFERSIZE,0);
    recv(sock,recv_mac ,20, 0);
    pipe.read((byte *)our_mac,20,1);
    if(memcmp(recv_mac,our_mac,20)) {
      std::cout << "WARNING: MAC in message failed to verify\n";
      exit(EXIT_FAILURE);
    }

    filestr.write(recv_buf, write_size);
    cum_send_size += write_size;
  }
  printf("Saved in all: %6d bytes\n",cum_send_size);
  filestr.flush();
  filestr.close();
  printf("\nFinish receiving file: %s\nFile saved as: %s\n",get_file,save_file);
  return 0;
}
